import logging
import ray

from calc_manager.base.runnable import Runnable
from typing import Optional
from calc_manager.base.task import Task, TaskFactory, DefaultRunnableTask, TaskInfo
from calc_manager.algorithm.algo import Algorithm
from calc_manager.operator import Operator
from calc_manager.plugins.plugins_manager import PluginManager
from typing import Dict, Any

logger = logging.getLogger(__name__)


@ray.remote
class RayRemoteRunnableTask(DefaultRunnableTask):
    logging.basicConfig(level=logging.DEBUG)

    def dispose(self):
        ray.actor.exit_actor()


# TODO: remove this
class RayRemoteRunnableTaskFatory(TaskFactory[RayRemoteRunnableTask]):
    def __init__(self, plugin_manager: Optional[PluginManager] = None, plugins_argv: Dict[str, Any] = {}):
        self.plugin_manager = plugin_manager
        self.plugins_argv = plugins_argv
    def create_task(self, info: TaskInfo) -> RayRemoteRunnableTask:
        task = RayRemoteRunnableTask.remote(info)

        def pre_run(task: RayRemoteRunnableTask):
            _ctx = task.get_ctx()
            if self.plugin_manager:
                invoke_param = Operator.get_plugin_objects_from_argv(self.plugin_manager, self.plugins_argv)
                task.invoke_param_raw = invoke_param
                invoke_param = {k: v.get_plugin_ojbect() for k, v in invoke_param.items()}
                task.runnable.set_extra_param(**invoke_param)
                logger.debug(f"pre_run invoke_param: {invoke_param}")   

        def post_run(task: RayRemoteRunnableTask, ret):
            ctx = task.get_ctx()
            logger.debug(f"post_run: {ctx}")
            # dispose plugins
            for k, v in task.invoke_param_raw.items():
                logger.debug(f"prepare to dispose plugin {k}")
                v.dispose()

        task.add_pre_run_callable.remote(pre_run)
        task.add_post_run_callable.remote(post_run)

        return task

class RayAlgoSingleTaskFactory(TaskFactory[RayRemoteRunnableTask]):
    def __init__(self, algo: Algorithm):
        self.algo = algo

    #TODO: handle ctx you can put ctx to task from pre_run an post_run
    def create_task(self, info: TaskInfo) -> RayRemoteRunnableTask:
        task = RayRemoteRunnableTask.remote(info)
        return task

#TODO: move pre_run and post run to bss level, and delete RayAlgoSingleTaskFactory
class RayAlgoOperatorTaskFactory(TaskFactory[RayRemoteRunnableTask]):
    def __init__(self, algo: Algorithm):
        self.algo = algo

    def create_task(self, info: TaskInfo) -> RayRemoteRunnableTask:
        task = RayRemoteRunnableTask.remote(info)

        def pre_run(task: RayRemoteRunnableTask):
            ctx = task.get_ctx()

            #TODO: optimize: avoid load operator again and again
            op_graph_id = task.runnable.runnable_name
            cur_op = self.algo.get_operator_by_graph_id(op_graph_id)

            argv = self.algo.algo_params
            #self.algo.mapping_extra_param_to_argv(
            #    ctx, task.runnable, argv
            #)
            self.algo.mapping_operator_plugins_to_extra_param(
                ctx, task.runnable, argv, cur_op
            )

            #pop out raw plugin info
            argv.pop(Operator.INVOKE_PLUGINS_KEY, None)

            self.algo.mapping_param_to_ctx_argv(
                ctx, task.runnable, argv
            )
            # Only print keys from ctx, including nested keys
            def get_nested_keys(d, prefix=''):
                keys = []
                for k, v in d.items():
                    new_key = f"{prefix}.{k}" if prefix else k
                    keys.append(new_key)
                    if isinstance(v, dict):
                        keys.extend(get_nested_keys(v, new_key))
                return keys
                
            ctx_keys = get_nested_keys(ctx)
            logger.debug(f"pre_run ctx keys: {ctx_keys}")

        def post_run(task: RayRemoteRunnableTask, ret):
            ctx = task.get_ctx()
            self.algo.mapping_result_to_ctx(ctx, task.runnable, ret)
            # Only print keys from ctx, including nested keys
            def get_nested_keys(d, prefix=''):
                keys = []
                for k, v in d.items():
                    new_key = f"{prefix}.{k}" if prefix else k
                    keys.append(new_key)
                    if isinstance(v, dict):
                        keys.extend(get_nested_keys(v, new_key))
                return keys
                
            ctx_keys = get_nested_keys(ctx)
            logger.debug(f"post_run ctx keys: {ctx_keys}")

        task.add_pre_run_callable.remote(pre_run)
        task.add_post_run_callable.remote(post_run)

        return task
